# import segmentation_models_pytorch as smp
import timm
import torch.nn as nn
from torch.cuda.amp import GradScaler, autocast

# class seg_builder(nn.Module):
    # def __init__(self, model_name, n_class):
    #     super().__init__()  
    #     self.model = smp.UnetPlusPlus(# UnetPlusPlus 
    #             encoder_name=model_name,        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    #             encoder_weights="imagenet",     # use `imagenet` pretrained weights for encoder initialization
    #             in_channels=3,                  # model input channels (1 for grayscale images, 3 for RGB, etc.)
    #             classes=n_class,                      # model output channels (number of classes in your dataset)
    #         )
    # @autocast()
    # def forward(self, x):
    #     #with autocast():
    #     x = self.model(x)
    #     return x

class timm_builder(nn.Module):
    def __init__(self, model_name, n_class, pretrained=True):
        super().__init__()  
        self.model = timm.create_model(
            model_name=model_name,
            num_classes=n_class,
            pretrained=pretrained,
        )
    @autocast()
    def forward(self, x):
        x = self.model(x)
        return x

